#include <src/pcps.h>
#include <src/blas_lapack.h>

// function to contract the vxI products with the contraction matrix and inverse
template<class T>
void vxIcontract(const int nap, const T * const ocm, const T * const inv, const T * const vxI, T * const output) {
  std::vector<T> work(nap, pcps::zero<T>());
  pcps::xgemm('N', 'T', 1, nap, nap, pcps::unity<T>(),         vxI, 1, ocm, nap, pcps::zero<T>(), &work.at(0), 1);
  pcps::xgemm('N', 'T', 1, nap, nap, pcps::unity<T>(), &work.at(0), 1, inv, nap, pcps::zero<T>(),      output, 1);
}

///  The cost of this function should scale asymptotically as (occupied)^2 * (virtual)
///
///  Assumptions:
///
///    Columns are labeled by alpha particles, rows by beta.
///

template<class T>
void agp_chol_der_contract_a(const int nas,
                             const int nap,
                             const T * const pairing_mat, // nas by nas
                             const T * const inverse_mat, // nap by nap
                             const T * const contraction_mat, // nas by nas
                             const T * const u2, // nap
                             const T * const u3, // nap
                             const T * const v2, // nap
                             const T * const v3, // nap
                             const int * const part_to_site_a, // size nap
                             const int * const part_to_site_b, // size nap
                             const int * const unocc_sites_a, // size ( nas - nap )
                             T * const output_tensor) // size nas by nas
{

  // zero the output tensor
  for (int i = 0; i < nas*nas; i++)
    output_tensor[i] = pcps::zero<T>();

  // get the number of unoccupied alpha sites
  const int nau = nas - nap;

  // get space for u1 vector
  std::vector<T> u1(nap, pcps::zero<T>());

  // get space for inverse times vectors (u = row, v = column)
  std::vector<T> Iu1(nap, pcps::zero<T>());
  std::vector<T> Iu2(nap, pcps::zero<T>());
  std::vector<T> Iu3(nap, pcps::zero<T>());
  std::vector<T> v1I(nap, pcps::zero<T>());
  std::vector<T> v2I(nap, pcps::zero<T>());
  std::vector<T> v3I(nap, pcps::zero<T>());

  // get space for the occupied contraction matrix
  std::vector<T> ocm(nap*nap, pcps::zero<T>());

  // populate the occupied contraction matrix
  for (int i = 0; i < nap; i++)
  for (int j = 0; j < nap; j++)
    ocm[nap*i+j] = contraction_mat[ nas * part_to_site_b[i] + part_to_site_a[j] ];

  // compute products of inverse matrix with noise vectors
  pcps::xgemm('T', 'N', nap, 1, nap, pcps::unity<T>(), inverse_mat, nap, u2, nap, pcps::zero<T>(), &Iu2.at(0), nap);
  pcps::xgemm('T', 'N', nap, 1, nap, pcps::unity<T>(), inverse_mat, nap, u3, nap, pcps::zero<T>(), &Iu3.at(0), nap);
  pcps::xgemm('N', 'T', 1, nap, nap, pcps::unity<T>(), v2, 1, inverse_mat, nap, pcps::zero<T>(), &v2I.at(0), 1);
  pcps::xgemm('N', 'T', 1, nap, nap, pcps::unity<T>(), v3, 1, inverse_mat, nap, pcps::zero<T>(), &v3I.at(0), 1);

  // get space for contractions of inverse times vectors with contraction matrix
  std::vector<T> v1Ic(nap, pcps::zero<T>());
  std::vector<T> v2Ic(nap, pcps::zero<T>());
  std::vector<T> v3Ic(nap, pcps::zero<T>());

  // get contractions for noise vectors
  vxIcontract(nap, &ocm.at(0), inverse_mat, &v2I.at(0), &v2Ic.at(0));
  vxIcontract(nap, &ocm.at(0), inverse_mat, &v3I.at(0), &v3Ic.at(0));

  // compute contraction of inverse matrix with the occupied contraction matrix
  T precontracted_inv = pcps::zero<T>();
  for (int i = 0; i < nap; i++)
  for (int j = 0; j < nap; j++)
    precontracted_inv += ocm[nap*i+j] * inverse_mat[nap*j+i];

  // loop over occupied alpha sites
  for (int col = 0; col < nap; col++) {

    // get the site this column corresponds to
    const int s = part_to_site_a[col];

    // construct the v1I product
    for (int i = 0; i < nap; i++)
      v1I[i] = inverse_mat[nap*col+i];

    // contract the v1I product with the contraction matrix
    vxIcontract(nap, &ocm.at(0), inverse_mat, &v1I.at(0), &v1Ic.at(0));

    // get inner products that don't depend on r
    const T ip12 = pcps::xdot(nap, &v1I.at(0), 1, u2, 1);
    const T ip13 = pcps::xdot(nap, &v1I.at(0), 1, u3, 1);
    const T ip22 = pcps::xdot(nap, &v2I.at(0), 1, u2, 1);
    const T ip23 = pcps::xdot(nap, &v2I.at(0), 1, u3, 1);
    const T ip32 = pcps::xdot(nap, &v3I.at(0), 1, u2, 1);
    const T ip33 = pcps::xdot(nap, &v3I.at(0), 1, u3, 1);

    // compute the full contractions that don't depend on r
    const T fc12 = pcps::xdot(nap, &v1Ic.at(0), 1, u2, 1);
    const T fc13 = pcps::xdot(nap, &v1Ic.at(0), 1, u3, 1);
    const T fc22 = pcps::xdot(nap, &v2Ic.at(0), 1, u2, 1);
    const T fc23 = pcps::xdot(nap, &v2Ic.at(0), 1, u3, 1);
    const T fc32 = pcps::xdot(nap, &v3Ic.at(0), 1, u2, 1);
    const T fc33 = pcps::xdot(nap, &v3Ic.at(0), 1, u3, 1);

    // loop over unoccupied alpha sites (and the diagonal site)
    for (int unocc = 0; unocc <= nau; unocc++) {

      // get the site index
      const int r = ( unocc < nau ? unocc_sites_a[unocc] : s );

      // construct the u1 vector, which is the difference between the new and old columns
      // of the pairing matrix
      for (int i = 0; i < nap; i++)
        u1[i] =   pairing_mat[ nas * part_to_site_b[i] + r ]
                - pairing_mat[ nas * part_to_site_b[i] + s ];

      // compute the remaining inner products
      const T ip11 = pcps::xdot(nap, &v1I[0], 1, &u1[0], 1);
      const T ip21 = pcps::xdot(nap, &v2I[0], 1, &u1[0], 1);
      const T ip31 = pcps::xdot(nap, &v3I[0], 1, &u1[0], 1);

      // compute the contribution from the precontracted inverse
      output_tensor[r*nas+s] += precontracted_inv * (   pcps::unity<T>()
                                                      + ip11 + ip22 + ip33
                                                      + ip11*ip22 + ip11*ip33 + ip22*ip33
                                                      - ip12*ip21 - ip13*ip31 - ip23*ip32
                                                      + ip11*ip22*ip33 - ip11*ip23*ip32
                                                      - ip12*ip21*ip33 - ip13*ip22*ip31
                                                      + ip12*ip23*ip31 + ip13*ip21*ip32 );

      // compute the remaining full contractions
      const T fc11 = pcps::xdot(nap, &v1Ic[0], 1, &u1[0], 1);
      const T fc21 = pcps::xdot(nap, &v2Ic[0], 1, &u1[0], 1);
      const T fc31 = pcps::xdot(nap, &v3Ic[0], 1, &u1[0], 1);

      // compute the contribution when v1 is factored to the left
      output_tensor[r*nas+s] -=   fc11 * ( pcps::unity<T>() + ip22 + ip33 + ip22*ip33 - ip32*ip23 )
                                + fc12 * (                  - ip21        + ip23*ip31 - ip21*ip33 )
                                + fc13 * (                  - ip31        + ip21*ip32 - ip31*ip22 );

      // compute the contribution when v2 is factored to the left
      output_tensor[r*nas+s] -=   fc22 * ( pcps::unity<T>() + ip11 + ip33 + ip11*ip33 - ip31*ip13 )
                                + fc21 * (                  - ip12        + ip13*ip32 - ip12*ip33 )
                                + fc23 * (                  - ip32        + ip12*ip31 - ip32*ip11 );

      // compute the contribution when v3 is factored to the left
      output_tensor[r*nas+s] -=   fc33 * ( pcps::unity<T>() + ip22 + ip11 + ip22*ip11 - ip12*ip21 )
                                + fc32 * (                  - ip23        + ip21*ip13 - ip23*ip11 )
                                + fc31 * (                  - ip13        + ip23*ip12 - ip13*ip22 );

      // construct the u1 contracting vector, which is the difference between the new and old columns
      // of the contraction matrix
      for (int i = 0; i < nap; i++)
        u1[i] =   contraction_mat[ nas * part_to_site_b[i] + r ]
                - contraction_mat[ nas * part_to_site_b[i] + s ];

      // compute the "u derivative" full contractions
      const T ud11 = pcps::xdot(nap, &v1I[0], 1, &u1[0], 1);
      const T ud21 = pcps::xdot(nap, &v2I[0], 1, &u1[0], 1);
      const T ud31 = pcps::xdot(nap, &v3I[0], 1, &u1[0], 1);

      // compute the "u derivative" contribution
      output_tensor[r*nas+s] +=   ud11 * ( pcps::unity<T>() + ip22 + ip33 + ip22*ip33 - ip32*ip23 )
                                + ud21 * (                  - ip12        + ip13*ip32 - ip12*ip33 )
                                + ud31 * (                  - ip13        + ip12*ip23 - ip13*ip22 );

    }

  }

}

template<class T>
void add_transpose(const int n, const T * const mat, T * const trn) {
  for (int i = 0; i < n; i++)
  for (int j = 0; j < n; j++)
    trn[j*n+i] += mat[i*n+j];
}

template<class T>
void agp_chol_der_contract(const int nas, // number of alpha sites
                           const int nap, // number of alpha particles
                           const T * const pairing_mat, // nsite/2 by nsite/2
                           const T * const inverse_mat, // npart/2 by npart/2
                           const T * const contrct_mat, // nsite/2 by nsite/2
                           const T * const u2, // npart/2
                           const T * const u3, // npart/2
                           const T * const v2, // npart/2
                           const T * const v3, // npart/2
                           const int * const part_to_site_a, // size npart/2
                           const int * const part_to_site_b, // size npart/2
                           const int * const unocc_sites_a, // size ( nsites/2 - npart/2 )
                           const int * const unocc_sites_b, // size ( nsites/2 - npart/2 )
                           T * const output_tensor_a, // size nsite/2 by nsite/2
                           T * const output_tensor_b) // size nsite/2 by nsite/2
{

  // alpha case
  agp_chol_der_contract_a(nas, nap, pairing_mat, inverse_mat, contrct_mat, u2, u3, v2, v3,
                          part_to_site_a, part_to_site_b, unocc_sites_a, output_tensor_a);

  // beta case (computed as the transpose of the alpha case)
  std::vector<T> pairing_trn(nas*nas, pcps::zero<T>());  add_transpose(nas, pairing_mat, &pairing_trn.at(0));
  std::vector<T> inverse_trn(nap*nap, pcps::zero<T>());  add_transpose(nap, inverse_mat, &inverse_trn.at(0));
  std::vector<T> contrct_trn(nas*nas, pcps::zero<T>());  add_transpose(nas, contrct_mat, &contrct_trn.at(0));
  agp_chol_der_contract_a(nas, nap, &pairing_trn.at(0), &inverse_trn.at(0), &contrct_trn.at(0), v2, v3, u2, u3,
                          part_to_site_b, part_to_site_a, unocc_sites_b, output_tensor_b);

}

template void agp_chol_der_contract(const int nas,
                                    const int nap,
                                    const double * const pairing_mat,
                                    const double * const inverse_mat,
                                    const double * const contrct_mat,
                                    const double * const u2,
                                    const double * const u3,
                                    const double * const v2,
                                    const double * const v3,
                                    const int * const part_to_site_a,
                                    const int * const part_to_site_b,
                                    const int * const unocc_sites_a,
                                    const int * const unocc_sites_b,
                                    double * const output_tensor_a,
                                    double * const output_tensor_b);
template void agp_chol_der_contract(const int nas,
                                    const int nap,
                                    const std::complex<double> * const pairing_mat,
                                    const std::complex<double> * const inverse_mat,
                                    const std::complex<double> * const contrct_mat,
                                    const std::complex<double> * const u2,
                                    const std::complex<double> * const u3,
                                    const std::complex<double> * const v2,
                                    const std::complex<double> * const v3,
                                    const int * const part_to_site_a,
                                    const int * const part_to_site_b,
                                    const int * const unocc_sites_a,
                                    const int * const unocc_sites_b,
                                    std::complex<double> * const output_tensor_a,
                                    std::complex<double> * const output_tensor_b);

template<class T>
void agp_chol_der_expand_a(const int nas,
                           const int nap,
                           const T * const pairing_mat, // nas by nas
                           const T * const inverse_mat, // nap by nap
                           const T * const contraction_mat, // nas by nas
                           const T * const u2, // nap
                           const T * const u3, // nap
                           const T * const v2, // nap
                           const T * const v3, // nap
                           const int * const part_to_site_a, // size nap
                           const int * const part_to_site_b, // size nap
                           const int * const unocc_sites_a, // size ( nas - nap )
                           T * const output_tensor) // size nas by nas
{

  // get the number of unoccupied alpha sites
  const int nau = nas - nap;

  // get space for u1 vector
  std::vector<T> u1(nap, pcps::zero<T>());

  // get space for inverse times vectors (u = row, v = column)
  std::vector<T> Iu1(nap, pcps::zero<T>());
  std::vector<T> Iu2(nap, pcps::zero<T>());
  std::vector<T> Iu3(nap, pcps::zero<T>());
  std::vector<T> v1I(nap, pcps::zero<T>());
  std::vector<T> v2I(nap, pcps::zero<T>());
  std::vector<T> v3I(nap, pcps::zero<T>());

  // initialize the occupied expansion matrix
  std::vector<T> oem(nap*nap, pcps::zero<T>());

  // prepare the occupied contraction matrix
  std::vector<T> ocm(nap*nau, pcps::zero<T>());
  for (int p = 0; p < nap; p++)
  for (int q = 0; q < nau; q++)
    ocm[nau*p+q] = contraction_mat[ nas * part_to_site_a[p] + unocc_sites_a[q] ];

  // compute products of inverse matrix with noise vectors
  pcps::xgemm('T', 'N', nap, 1, nap, pcps::unity<T>(), inverse_mat, nap, u2, nap, pcps::zero<T>(), &Iu2.at(0), nap);
  pcps::xgemm('T', 'N', nap, 1, nap, pcps::unity<T>(), inverse_mat, nap, u3, nap, pcps::zero<T>(), &Iu3.at(0), nap);
  pcps::xgemm('N', 'T', 1, nap, nap, pcps::unity<T>(), v2, 1, inverse_mat, nap, pcps::zero<T>(), &v2I.at(0), 1);
  pcps::xgemm('N', 'T', 1, nap, nap, pcps::unity<T>(), v3, 1, inverse_mat, nap, pcps::zero<T>(), &v3I.at(0), 1);

  // initialize the constant used in the "easy term"
  T cont_det_ratios = pcps::zero<T>();

  // initialize the sum of contraction matrix elements
  T cme_sum = pcps::zero<T>();

  // initialize variables for inner product contractions
  T ipc11 = pcps::zero<T>();
  T ipc12 = pcps::zero<T>();
  T ipc13 = pcps::zero<T>();
  T ipc21 = pcps::zero<T>();
  T ipc22 = pcps::zero<T>();
  T ipc23 = pcps::zero<T>();
  T ipc31 = pcps::zero<T>();
  T ipc32 = pcps::zero<T>();
  T ipc33 = pcps::zero<T>();
  T ipc_11_22 = pcps::zero<T>();
  T ipc_11_33 = pcps::zero<T>();
  T ipc_11_23 = pcps::zero<T>();
  T ipc_11_32 = pcps::zero<T>();
  T ipc_12_21 = pcps::zero<T>();
  T ipc_13_31 = pcps::zero<T>();
  T ipc_12_31 = pcps::zero<T>();
  T ipc_13_21 = pcps::zero<T>();

  // initialize variables for partial sums of inner product contractions
  //std::vector<T> ipc11_psums(nap, pcps::zero<T>());
  //std::vector<T> ipc12_psums(nap, pcps::zero<T>());
  //std::vector<T> ipc13_psums(nap, pcps::zero<T>());
  std::vector<T> ipc21_psums(nap, pcps::zero<T>());
  //std::vector<T> ipc22_psums(nap, pcps::zero<T>());
  //std::vector<T> ipc23_psums(nap, pcps::zero<T>());
  std::vector<T> ipc31_psums(nap, pcps::zero<T>());
  //std::vector<T> ipc32_psums(nap, pcps::zero<T>());
  //std::vector<T> ipc33_psums(nap, pcps::zero<T>());
  //std::vector<T> ipc_11_22_psums(nap, pcps::zero<T>());
  //std::vector<T> ipc_11_33_psums(nap, pcps::zero<T>());
  //std::vector<T> ipc_11_23_psums(nap, pcps::zero<T>());
  //std::vector<T> ipc_11_32_psums(nap, pcps::zero<T>());
  //std::vector<T> ipc_12_21_psums(nap, pcps::zero<T>());
  //std::vector<T> ipc_13_31_psums(nap, pcps::zero<T>());
  //std::vector<T> ipc_12_31_psums(nap, pcps::zero<T>());
  //std::vector<T> ipc_13_21_psums(nap, pcps::zero<T>());

  // compute the inner products that don't depend on p or q
  const T ip22 = pcps::xdot(nap, &v2I.at(0), 1, u2, 1);
  const T ip23 = pcps::xdot(nap, &v2I.at(0), 1, u3, 1);
  const T ip32 = pcps::xdot(nap, &v3I.at(0), 1, u2, 1);
  const T ip33 = pcps::xdot(nap, &v3I.at(0), 1, u3, 1);

  // Fp is just the inverse matrix...
  const T * const Fp = inverse_mat;
  //std::vector<T> Fp(nap*nap, pcps::zero<T>());

  // loop over occupied alpha sites to compute Gp
  std::vector<T> Gp(nap*nap, pcps::zero<T>());
  for (int col = 0; col < nap; col++) {

    // get the site this column corresponds to
    const int p = part_to_site_a[col];

    // construct the p half of the u1 vector
    for (int i = 0; i < nap; i++)
      u1[i] = pairing_mat[ nas * part_to_site_b[i] + p ];

    // compute a row of Gp
    pcps::xgemm('T', 'N', nap, 1, nap, pcps::unity<T>(), inverse_mat, nap, &u1.at(0), nap, pcps::zero<T>(), &Gp.at(col*nap), nap);

  }

  // loop over unoccupied alpha sites to compute Gq
  std::vector<T> Gq(nau*nap, pcps::zero<T>());
  for (int unocc = 0; unocc < nau; unocc++) {

    // get the site index
    const int q = unocc_sites_a[unocc];

    // construct the q half of the u1 vector
    for (int i = 0; i < nap; i++)
      u1[i] = pairing_mat[ nas * part_to_site_b[i] + q ];

    // compute a row of Gq
    pcps::xgemm('T', 'N', nap, 1, nap, pcps::unity<T>(), inverse_mat, nap, &u1.at(0), nap, pcps::zero<T>(), &Gq.at(unocc*nap), nap);

  }

  // initialize a matrix to hold the contraction of Gq and Gp with the occupied contraction matrix
  std::vector<T> Gc(nap*nap, pcps::zero<T>());

  // initialize vectors that will hold contractions of inner products and Gp and Gq
  std::vector<T> G12(nap, pcps::zero<T>());
  std::vector<T> G13(nap, pcps::zero<T>());

  // loop over occupied alpha sites
  for (int col = 0; col < nap; col++) {

    // get the site this column corresponds to
    const int p = part_to_site_a[col];

    // construct the v1I product
    for (int i = 0; i < nap; i++)
      v1I[i] = inverse_mat[nap*col+i];

    // get inner products that don't depend on q
    const T ip12 = pcps::xdot(nap, &v1I.at(0), 1, u2, 1);
    const T ip13 = pcps::xdot(nap, &v1I.at(0), 1, u3, 1);

    // loop over unoccupied alpha sites (and the diagonal site)
    for (int unocc = 0; unocc <= nau; unocc++) {

      // get the site index
      const int q = ( unocc < nau ? unocc_sites_a[unocc] : p );

      // construct the u1 vector, which is the difference between the new and old columns
      // of the pairing matrix
      for (int i = 0; i < nap; i++)
        u1[i] =   pairing_mat[ nas * part_to_site_b[i] + q ]
                - pairing_mat[ nas * part_to_site_b[i] + p ];

      // compute the remaining inner products
      const T ip11 = pcps::xdot(nap, &v1I[0], 1, &u1[0], 1);
      const T ip21 = pcps::xdot(nap, &v2I[0], 1, &u1[0], 1);
      const T ip31 = pcps::xdot(nap, &v3I[0], 1, &u1[0], 1);

      // get the pq element of the contraction matrix
      const T cont_elem = contraction_mat[p*nas+q];

      // compute the contribution to the contraction of the determinant ratios
      cont_det_ratios += cont_elem * (   pcps::unity<T>()
                                       + ip11 + ip22 + ip33
                                       + ip11*ip22 + ip11*ip33 + ip22*ip33
                                       - ip12*ip21 - ip13*ip31 - ip23*ip32
                                       + ip11*ip22*ip33 - ip11*ip23*ip32
                                       - ip12*ip21*ip33 - ip13*ip22*ip31
                                       + ip12*ip23*ip31 + ip13*ip21*ip32 );

      // add the contraction matrix element to the sum of contraction elements
      cme_sum += cont_elem;

      // add the contributions to the contraction of inner products
      ipc11 += cont_elem * ip11;
      ipc12 += cont_elem * ip12;
      ipc13 += cont_elem * ip13;
      ipc21 += cont_elem * ip21;
      ipc22 += cont_elem * ip22;
      ipc23 += cont_elem * ip23;
      ipc31 += cont_elem * ip31;
      ipc32 += cont_elem * ip32;
      ipc33 += cont_elem * ip33;
      ipc_11_22 += cont_elem * ip11 * ip22;
      ipc_11_33 += cont_elem * ip11 * ip33;
      ipc_11_23 += cont_elem * ip11 * ip23;
      ipc_11_32 += cont_elem * ip11 * ip32;
      ipc_12_21 += cont_elem * ip12 * ip21;
      ipc_13_31 += cont_elem * ip13 * ip31;
      ipc_12_31 += cont_elem * ip12 * ip31;
      ipc_13_21 += cont_elem * ip13 * ip21;

      // add the contributions to the partial sums of contractions of inner products
      //ipc11_psums[col] += cont_elem * ip11;
      //ipc12_psums[col] += cont_elem * ip12;
      //ipc13_psums[col] += cont_elem * ip13;
      ipc21_psums[col] += cont_elem * ip21;
      //ipc22_psums[col] += cont_elem * ip22;
      //ipc23_psums[col] += cont_elem * ip23;
      ipc31_psums[col] += cont_elem * ip31;
      //ipc32_psums[col] += cont_elem * ip32;
      //ipc33_psums[col] += cont_elem * ip33;
      //ipc_11_22_psums[col] += cont_elem * ip11 * ip22;
      //ipc_11_33_psums[col] += cont_elem * ip11 * ip33;
      //ipc_11_23_psums[col] += cont_elem * ip11 * ip23;
      //ipc_11_32_psums[col] += cont_elem * ip11 * ip32;
      //ipc_12_21_psums[col] += cont_elem * ip12 * ip21;
      //ipc_13_31_psums[col] += cont_elem * ip13 * ip31;
      //ipc_12_31_psums[col] += cont_elem * ip12 * ip31;
      //ipc_13_21_psums[col] += cont_elem * ip13 * ip21;

    }

    // compute the partial sum of contraction elements, skipping the p == q term
    T cme_psum = pcps::zero<T>();
    for (int unocc = 0; unocc < nau; unocc++)
      cme_psum += ocm[col*nau+unocc];

    // compute Gq's contribution to this row of the Gc matrix
    pcps::xgemm('N', 'N', nap, 1, nau, pcps::unity<T>(), &Gq.at(0), nap, &ocm.at(col*nau), nau, pcps::zero<T>(), &Gc.at(col*nap), nap);

    // compute Gp's contribution to this row of the Gc matrix
    pcps::xaxpy(nap, -cme_psum, &Gp.at(col*nap), 1, &Gc.at(col*nap), 1);

    // compute the contributions to the contractions of inner products with Gc
    pcps::xaxpy(nap, ip12, &Gc.at(col*nap), 1, &G12.at(0), 1);
    pcps::xaxpy(nap, ip13, &Gc.at(col*nap), 1, &G13.at(0), 1);

    // compute the u1 derivative's green "q" contribution
    {
      const T * const ocm_ptr = &ocm.at(col*nau);
      const T x = pcps::unity<T>() + ip22 + ip33 + ip22*ip33 - ip23*ip32;
      for (int i = 0; i < nap; i++) {
        T * const out_ptr = output_tensor + nas * part_to_site_b[i];
        const T d = x * v1I[i];
        for (int unocc = 0; unocc < nau; unocc++)
          out_ptr[unocc_sites_a[unocc]] += d * ocm_ptr[unocc];
      }
    }

    // compute the u1 derivative's green "p" contribution
    {
      const T x = pcps::unity<T>() + ip22 + ip33 + ip22*ip33 - ip23*ip32;
      for (int i = 0; i < nap; i++)
        output_tensor[ nas * part_to_site_b[i] + p ] -= x * v1I[i] * cme_psum;
    }

    // compute the u1 derivative's blue "21q" contribution
    {
      const T * const ocm_ptr = &ocm.at(col*nau);
      const T x = - ip12 - ip12*ip33 + ip13*ip32;
      for (int i = 0; i < nap; i++) {
        T * const out_ptr = output_tensor + nas * part_to_site_b[i];
        const T d = x * v2I[i];
        for (int unocc = 0; unocc < nau; unocc++)
          out_ptr[unocc_sites_a[unocc]] += d * ocm_ptr[unocc];
      }
    }

    // compute the u1 derivative's blue "31q" contribution
    {
      const T * const ocm_ptr = &ocm.at(col*nau);
      const T x = - ip13 - ip13*ip22 + ip12*ip23;
      for (int i = 0; i < nap; i++) {
        T * const out_ptr = output_tensor + nas * part_to_site_b[i];
        const T d = x * v3I[i];
        for (int unocc = 0; unocc < nau; unocc++)
          out_ptr[unocc_sites_a[unocc]] += d * ocm_ptr[unocc];
      }
    }

    // compute the u1 derivative's blue "21p" contribution
    {
      const T x = - ip12 - ip12*ip33 + ip13*ip32;
      for (int i = 0; i < nap; i++)
        output_tensor[ nas * part_to_site_b[i] + p ] -= x * v2I[i] * cme_psum;
    }

    // compute the u1 derivative's blue "31p" contribution
    {
      const T x = - ip13 - ip13*ip22 + ip12*ip23;
      for (int i = 0; i < nap; i++)
        output_tensor[ nas * part_to_site_b[i] + p ] -= x * v3I[i] * cme_psum;
    }

  }

  // add the "easy term"
  for (int i = 0; i < nap; i++)
  for (int j = 0; j < nap; i++)
    oem[nap*i+j] += inverse_mat[nap*j+i] * cont_det_ratios;

  // add the "22" green terms
  {
    const T x = pcps::unity<T>() + ipc11 + ipc33 + ipc_11_33 - ipc_13_31;
    for (int i = 0; i < nap; i++)
    for (int j = 0; j < nap; i++)
      oem[nap*i+j] -= v2I[i] * Iu2[j] * x;
  }

  // add the "33" green terms
  {
    const T x = pcps::unity<T>() + ipc11 + ipc22 + ipc_11_22 - ipc_12_21;
    for (int i = 0; i < nap; i++)
    for (int j = 0; j < nap; i++)
      oem[nap*i+j] -= v3I[i] * Iu3[j] * x;
  }

  // add the "23" green terms
  {
    const T x =                  - ipc32         - ipc_11_32 + ipc_12_31;
    for (int i = 0; i < nap; i++)
    for (int j = 0; j < nap; i++)
      oem[nap*i+j] -= v2I[i] * Iu3[j] * x;
  }

  // add the "32" green terms
  {
    const T x =                  - ipc23         - ipc_11_23 + ipc_13_21;
    for (int i = 0; i < nap; i++)
    for (int j = 0; j < nap; i++)
      oem[nap*i+j] -= v3I[i] * Iu2[j] * x;
  }

  // compute blue term
  {
    const T x = pcps::unity<T>() + ip22 + ip33 + ip22*ip33 - ip32*ip23;
    pcps::xgemm('N', 'T', nap, nap, nap, -x, &Gc.at(0), nap, Fp, nap, pcps::unity<T>(), &oem.at(0), nap);
  }

  // compute purple terms
  {

    // compute the contractions of Fp with the partial sums of the 21 and 31 inner product contractions
    std::vector<T> F21(nap, pcps::zero<T>());
    std::vector<T> F31(nap, pcps::zero<T>());
    pcps::xgemm('N', 'N', nap, 1, nap, pcps::unity<T>(), Fp, nap, &ipc21_psums.at(0), nap, pcps::zero<T>(), &F21.at(0), nap);
    pcps::xgemm('N', 'N', nap, 1, nap, pcps::unity<T>(), Fp, nap, &ipc31_psums.at(0), nap, pcps::zero<T>(), &F31.at(0), nap);

    // compute the "12" purple terms' contributions
    pcps::xgemm('N', 'N', nap, nap, 1, pcps::unity<T>() + ip33, &Iu2.at(0), nap, &F21.at(0), 1, pcps::unity<T>(), &oem.at(0), nap);
    pcps::xgemm('N', 'N', nap, nap, 1,                  - ip23, &Iu2.at(0), nap, &F31.at(0), 1, pcps::unity<T>(), &oem.at(0), nap);

    // compute the "13" purple terms' contributions
    pcps::xgemm('N', 'N', nap, nap, 1, pcps::unity<T>() + ip22, &Iu3.at(0), nap, &F31.at(0), 1, pcps::unity<T>(), &oem.at(0), nap);
    pcps::xgemm('N', 'N', nap, nap, 1,                  - ip32, &Iu3.at(0), nap, &F21.at(0), 1, pcps::unity<T>(), &oem.at(0), nap);

  }

  // compute yellow terms
  {

    // compute the "21" yellow terms' contributions
    pcps::xgemm('N', 'N', nap, nap, 1, pcps::unity<T>() + ip33, &G12.at(0), nap, &v2I.at(0), 1, pcps::unity<T>(), &oem.at(0), nap);
    pcps::xgemm('N', 'N', nap, nap, 1,                  - ip32, &G13.at(0), nap, &v2I.at(0), 1, pcps::unity<T>(), &oem.at(0), nap);

    // compute the "31" yellow terms' contributions
    pcps::xgemm('N', 'N', nap, nap, 1, pcps::unity<T>() + ip22, &G13.at(0), nap, &v3I.at(0), 1, pcps::unity<T>(), &oem.at(0), nap);
    pcps::xgemm('N', 'N', nap, nap, 1,                  - ip23, &G12.at(0), nap, &v3I.at(0), 1, pcps::unity<T>(), &oem.at(0), nap);

  }

  // add occupied expansion matrix to the output tensor
  for (int i = 0; i < nap; i++)
  for (int j = 0; j < nap; i++)
    output_tensor[ nap * part_to_site_b[i] + part_to_site_a[j] ] += oem[nap*i+j];

}

template<class T>
void agp_chol_der_expand(const int nas, // number of alpha sites
                         const int nap, // number of alpha particles
                         const T * const pairing_mat, // nsite/2 by nsite/2
                         const T * const inverse_mat, // npart/2 by npart/2
                         const T * const contrct_mat_a, // nsite/2 by nsite/2
                         const T * const contrct_mat_b, // nsite/2 by nsite/2
                         const T * const u2, // npart/2
                         const T * const u3, // npart/2
                         const T * const v2, // npart/2
                         const T * const v3, // npart/2
                         const int * const part_to_site_a, // size npart/2
                         const int * const part_to_site_b, // size npart/2
                         const int * const unocc_sites_a, // size ( nsites/2 - npart/2 )
                         const int * const unocc_sites_b, // size ( nsites/2 - npart/2 )
                         T * const output_tensor) // size nsite/2 by nsite/2
{

  // zero the output tensor
  for (int i = 0; i < nas*nas; i++)
    output_tensor[i] = pcps::zero<T>();

  // alpha case
  agp_chol_der_expand_a(nas, nap, pairing_mat, inverse_mat, contrct_mat_a, u2, u3, v2, v3,
                        part_to_site_a, part_to_site_b, unocc_sites_a, output_tensor);

  // beta case (computed as the transpose of the alpha case)
  std::vector<T> pairing_trn(nas*nas, pcps::zero<T>());  add_transpose(nas, pairing_mat, &pairing_trn.at(0));
  std::vector<T> inverse_trn(nap*nap, pcps::zero<T>());  add_transpose(nap, inverse_mat, &inverse_trn.at(0));
  std::vector<T>  output_trn(nas*nas, pcps::zero<T>());
  agp_chol_der_expand_a(nas, nap, &pairing_trn.at(0), &inverse_trn.at(0), contrct_mat_b, v2, v3, u2, u3,
                        part_to_site_b, part_to_site_a, unocc_sites_b, &output_trn.at(0));
  add_transpose(nas, &output_trn.at(0), output_tensor);

}

template void agp_chol_der_expand(const int nas,
                                    const int nap,
                                    const double * const pairing_mat,
                                    const double * const inverse_mat,
                                    const double * const contrct_mat_a,
                                    const double * const contrct_mat_b,
                                    const double * const u2,
                                    const double * const u3,
                                    const double * const v2,
                                    const double * const v3,
                                    const int * const part_to_site_a,
                                    const int * const part_to_site_b,
                                    const int * const unocc_sites_a,
                                    const int * const unocc_sites_b,
                                    double * const output_tensor);
template void agp_chol_der_expand(const int nas,
                                  const int nap,
                                  const std::complex<double> * const pairing_mat,
                                  const std::complex<double> * const inverse_mat,
                                  const std::complex<double> * const contrct_mat_a,
                                  const std::complex<double> * const contrct_mat_b,
                                  const std::complex<double> * const u2,
                                  const std::complex<double> * const u3,
                                  const std::complex<double> * const v2,
                                  const std::complex<double> * const v3,
                                  const int * const part_to_site_a,
                                  const int * const part_to_site_b,
                                  const int * const unocc_sites_a,
                                  const int * const unocc_sites_b,
                                  std::complex<double> * const output_tensor);

